#coding=utf-8
import argparse
import os 
import xml.dom.minidom as minidom
from utils import visualization_utils as vis_util
from utils import label_map_util
import numpy as np
import tensorflow as tf
from matplotlib import pyplot as plt
from PIL import Image
import xml.etree.ElementTree as ET
import re
import classify_image2 as cls
import cv2
NUM_CLASSES = 90
model_dir='C:\\code'
output_dir='C:\\code'
image_dir='C:\\code\\org'
'''
def parse_args(check=True):
    parser = argparse.ArgumentParser()
    #模型pb文件labletxt文件所在目录
    parser.add_argument('--model_dir', type=str, required=True)
    #xml输出目录
    parser.add_argument('--output_dir', type=str, required=True)
    #图片所在目录----要求图片名以  分类名+'@'+文件名+'.jpg' 命名
    parser.add_argument('--image_dir', type=str, required=True)
    FLAGS, unparsed = parser.parse_known_args()
    return FLAGS, unparsed

def prettyXml(element, indent, newline, level = 0): # elemnt为传进来的Elment类，参数indent用于缩进，newline用于换行  
    if element:  # 判断element是否有子元素  
        if element.text == None or element.text.isspace(): # 如果element的text没有内容  
            element.text = newline + indent * (level + 1)    
        else:  
            element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)  
    #else:  # 此处两行如果把注释去掉，Element的text也会另起一行  
        #element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level  
    temp = list(element) # 将elemnt转成list  
    for subelement in temp:  
        if temp.index(subelement) < (len(temp) - 1): # 如果不是list的最后一个元素，说明下一个行是同级别元素的起始，缩进应一致  
            subelement.tail = newline + indent * (level + 1)  
        else:  # 如果是list的最后一个元素， 说明下一行是母元素的结束，缩进应该少一个  
            subelement.tail = newline + indent * level  
'''


    #print('1')
PATH_TO_CKPT = os.path.join(model_dir, 'ssd_inception_v2.pb')
PATH_TO_LABELS = os.path.join(model_dir, 'mscoco_label_map.pbtxt')

#print('2')

detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')
#print('3')
label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)
#print('4')
def load_image_into_numpy_array(image):
    (im_width, im_height) = image.size
    return np.array(image.getdata()).reshape(
        (im_height, im_width, 3)).astype(np.uint8)


    #test_img_path = PATH_TO_IMAGE
    #print('5')
def cutphoto(image,xmin1,ymin1,xmax1,yamx1):
    xmin1 = np.ceil(xmin_l * im_width).astype(np.int32) #转成int类型
    ymin1 = np.ceil(ymin_l * im_height).astype(np.int32)
    xmax1 = np.ceil(xmax_l * im_width).astype(np.int32)
    ymax1 = np.ceil(ymax_l * im_height).astype(np.int32)
    image_c = image.crop((xmin1,ymin1,xmax1,ymax1)) # 根据box裁剪
    return (image_c)
with detection_graph.as_default():
    with tf.Session(graph=detection_graph) as sess:
        # print('6')
        #print(PATH_TO_IMAGE)

        list = str.split('')
        for aroot, dirs, files in os.walk(image_dir):
            list = files
            #print('list', list)
        for l in list:
            # image = './test3/' + i
            #image_data = open(l, 'rb').read()
#            print('l', l)
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
            detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
            detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
            detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
            num_detections = detection_graph.get_tensor_by_name('num_detections:0')
            image = Image.open(os.path.join(image_dir,l))
            image_np = load_image_into_numpy_array(image)
            image_np_expanded = np.expand_dims(image_np, axis=0)
            box_new = np.zeros(4)
            im_width,im_height = image.size
            '''
            plt.figure("image")
            plt.imshow(image)
            plt.show()
            '''
            box1=()
            (boxes, scores, classes, num) = sess.run(
                [detection_boxes, detection_scores, detection_classes, num_detections],
                feed_dict={image_tensor: image_np_expanded})
            box_new,box_to_color_map=vis_util.visualize_boxes_and_labels_on_image_array(
                                                                           image_np,
                                                                           np.squeeze(boxes),
                                                                           np.squeeze(classes).astype(np.int32),
                                                                           np.squeeze(scores),
                                                                           category_index,
                                                                           use_normalized_coordinates=True,
                                                                           line_thickness=8)   
            #a=box_new.shape[0]
            #print(box_new)
            if not box_new:
                print('no car!!!')
            else:
                i=0
                #print(box)
                for box,color in box_to_color_map.items():              
                    ymin_l, xmin_l, ymax_l, xmax_l = box
                    #print(ymin_l)
                    #print(box)
                    #print(box)
                    #print('ymin_l',ymin_l)
                    image_c=cutphoto(image,xmin_l,ymin_l,xmax_l,ymax_l)#调用裁切图片方法
                        #print(image_c)
        #                     xmin1 = np.ceil(xmin_l * im_width).astype(np.int32) #转成int类型
        #                     ymin1 = np.ceil(ymin_l * im_height).astype(np.int32)
        #                     xmax1 = np.ceil(xmax_l * im_width).astype(np.int32)
        #                     ymax1 = np.ceil(ymax_l * im_height).astype(np.int32)
        #                     boxw=xmax1-xmin1
        #                     boxh=ymax1-ymin1
        #                     area=boxw*boxh
        #                     print(area)
        #                     dropdata=area/(im_width*im_height)
        #                     print(dropdata)
                    photo='C:\\code\\tmp\\'+str(i)+'.jpg'#缓存裁切图片路径
                    #print(photo)
                    plt.imsave(photo,image_c)
                    #photo='C:\\code\\org\\'+l
                    labelname=cls.run_inference_on_image(photo,'C:\\code\\pb\\freezed.label','C:\\code\\pb\\freezed.pb')#调用测试获取label
                    #其中photo送入的裁切好的图片
                    #print(score)
                    print(labelname)
                    os.remove(photo)#删除缓存图片
                    display_str_list=vis_util.draw_bounding_box_on_image_array(
                                                                image_np,
                                                                ymin_l,
                                                                xmin_l,
                                                                ymax_l,
                                                                xmax_l,
                                                                color=color,
                                                                thickness=4,
                                                                display_str_list=labelname,
                                                                use_normalized_coordinates=True)#image_np要画框的图 box依次往图上画框
                        #print(display_str_list)
                path=os.path.join(output_dir,l)#出结果图路径
                plt.imsave(path,image_np)
                print('ok')
        #                     img = cv2.imread(photo)
        #                     colors = (0,0,255)
        #                     img=cv2.rectangle(img,(xmin1,ymin1),(xmax1,ymax1),colors, 5)
        #                     img=cv2.putText(img,s1,(xmin1,ymin1), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 0), 3)
        #                     i=i+1
        #             cv2.imwrite('C:\\code\\1221.jpg',img)
        #                         image_c = image.crop((xmin1,ymin1,xmax1,ymax1)) # 根据box裁剪
        #                         # prettyXml(root, '\t', '\n')
        #                         # xmlname=b_name+'.xml'
        #                         path = 'C:\\code\\cut'
        #                         print(ymax1)
        #                         # print(xmlname)
        #                         #indent(root)
        #                         #tree = ET.ElementTree(root)
        #                         #tree.write(os.path.join(PATH_TO_OUT, l_new), encoding="utf-8", xml_declaration=None)
        #                         plt.imsave(os.path.join(path,l), image_c)
        #                         break
                        